# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/15 12:20 
# @Author  : shutao
# @FileName: triplet_loss.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import paddle
import paddle.nn as nn
import paddle.nn.functional as F

from paddle import ParamAttr
from paddle.regularizer import L2Decay

from ..base_net import BaseNet


class TripletLoss(BaseNet):
    """Triplet loss function"""

    def __init__(self):
        super(TripletLoss, self).__init__()

        self.alpha = 0.2

    def forward(self, preds, labels):
        """
        triplet loss
        Args:
            preds: 预测生成的特征向量
            labels: 标签

        ..math:
            Loss = Count(||f(xa_i) - f(xp_i)||^2 - ||f(xa_i) - f(xn_i)||^2 + alpha)

        Returns:
            loss
        """
        assert len(preds) == 2, "输入 preds 的数量必须为2"
        assert type(labels) == paddle.Tensor and labels.shape[0] % 2 == 0, "输入 labels 的数量必须为2的倍数"
        imgs1, imgs2 = preds

        loss = 0.
        for idx, (img1, img2, lab) in enumerate(zip(imgs1, imgs2, labels)):
            diff = img1 - img2  # 计算向量欧氏距离
            l2_norm = diff.square().sum().sqrt()  # 计算 L2 范数

            if lab.numpy()[0] == 1:
                loss += l2_norm
            elif lab.numpy()[0] == 0:
                loss -= l2_norm
                loss += self.alpha
            else:
                raise ValueError("lab 参数错误，当前为：{}".format(lab.numpy()[0]))

        return loss
